#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <netinet/ip_icmp.h>
#include <time.h>

#define PACKET_SIZE 1024
#define ICMP_ECHO_REQUEST 8
#define TIMEOUT_SECONDS 2

// 计算校验和
unsigned short calculate_checksum(unsigned short *buf, int nwords) {
    unsigned long sum;
    for (sum = 0; nwords > 0; nwords--)
        sum += *buf++;
    while (sum >> 16)
        sum = (sum & 0xFFFF) + (sum >> 16);
    return ~sum;
}

// 发送ICMP回显请求
void send_icmp_request(int sockfd, struct sockaddr_in *dest_addr) {
    struct icmp icmp_msg;
    icmp_msg.icmp_type = ICMP_ECHO_REQUEST;
    icmp_msg.icmp_code = 0;
    icmp_msg.icmp_cksum = 0;
    icmp_msg.icmp_seq = 1;
    icmp_msg.icmp_id = getpid() & 0xFFFF;
    // 填充数据部分，这里简单填充为时间戳
    time_t t;
    time(&t);
    memcpy(icmp_msg.icmp_data, &t, sizeof(t));
    icmp_msg.icmp_cksum = calculate_checksum((unsigned short *)&icmp_msg, sizeof(icmp_msg) / sizeof(unsigned short));

    sendto(sockfd, &icmp_msg, sizeof(icmp_msg), MSG_CONFIRM, (const struct sockaddr *)dest_addr, sizeof(*dest_addr));
}

// 接收ICMP回显应答
int receive_icmp_reply(int sockfd, struct sockaddr_in *from_addr) {
    char buffer[PACKET_SIZE];
    socklen_t from_len = sizeof(*from_addr);
    struct ip *ip_header;
    struct icmp *icmp_header;
    int nbytes;

    fd_set read_fds;
    struct timeval timeout;
    FD_ZERO(&read_fds);
    FD_SET(sockfd, &read_fds);
    timeout.tv_sec = TIMEOUT_SECONDS;
    timeout.tv_usec = 0;

    int activity = select(sockfd + 1, &read_fds, NULL, NULL, &timeout);
    if (activity < 0) {
        perror("select error");
        return -1;
    } else if (activity == 0) {
        printf("Request timed out\n");
        return -1;
    }

    nbytes = recvfrom(sockfd, buffer, PACKET_SIZE, MSG_WAITALL, (struct sockaddr *)from_addr, &from_len);
    if (nbytes < 0) {
        perror("recvfrom error");
        return -1;
    }

    ip_header = (struct ip *)buffer;
    icmp_header = (struct icmp *)(buffer + (ip_header->ip_hl << 2));
    if (icmp_header->icmp_type == ICMP_ECHO_REPLY && icmp_header->icmp_id == (getpid() & 0xFFFF)) {
        time_t received_time;
        memcpy(&received_time, icmp_header->icmp_data, sizeof(received_time));
        printf("Received ICMP reply from %s: ", inet_ntoa(from_addr->sin_addr));
        printf("Time: %ld\n", (long)(time(NULL) - received_time));
        return 0;
    }
    return -1;
}

int main(int argc, char *argv[]) {
    if (argc != 2) {
        printf("Usage: %s <destination_ip>\n", argv[0]);
        return 1;
    }

    int sockfd = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
    if (sockfd < 0) {
        perror("socket creation failed");
        return 1;
    }

    struct sockaddr_in dest_addr;
    memset(&dest_addr, 0, sizeof(dest_addr));
    dest_addr.sin_family = AF_INET;
    dest_addr.sin_port = 0;
    if (inet_pton(AF_INET, argv[1], &dest_addr.sin_addr) <= 0) {
        perror("Invalid address/ Address not supported");
        close(sockfd);
        return 1;
    }

    send_icmp_request(sockfd, &dest_addr);
    if (receive_icmp_reply(sockfd, &dest_addr)!= 0) {
        printf("Failed to receive valid ICMP reply\n");
    }

    close(sockfd);
    return 0;
}